package cn.camel.algorithm.tree.rbt;

/**
 * @author: Camel
 * @description: 红黑树
 * @date: 2025/10/18 14:47
 */
public class RedBlackTree <T extends Comparable<T>>{
    private RBNode<T> root;
    private final RBNode<T> NIL; // 哨兵节点（叶子节点）

    public RedBlackTree() {
        NIL = new RBNode<>(null);
        NIL.setColor(Color.BLACK);
        root = NIL;
    }
    // 插入方法
    public void insert(T data) {
        RBNode<T> node = new RBNode<>(data);
        node.left = NIL;
        node.right = NIL;
        node.parent = NIL;

        RBNode<T> current = root;
        RBNode<T> parent = NIL;

        // 二叉搜索树插入
        while (current != NIL) {
            parent = current;
            if (node.getData().compareTo(current.getData()) < 0) {
                current = current.left;
            } else {
                current = current.right;
            }
        }

        node.parent = parent;
        if (parent == NIL) {
            root = node;
        } else if (node.getData().compareTo(parent.getData()) < 0) {
            parent.left = node;
        } else {
            parent.right = node;
        }

        fixInsert(node); // 修复红黑树性质
    }

    // 修复插入后的红黑树性质
    private void fixInsert(RBNode<T> node) {
        while (node.parent.getColor() == Color.RED) {
            if (node.parent == node.parent.parent.left) {
                RBNode<T> uncle = node.parent.parent.right;

                // Case 1: 叔叔节点是红色
                if (uncle.getColor() == Color.RED) {
                    node.parent.setColor(Color.BLACK);
                    uncle.setColor(Color.BLACK) ;
                    node.parent.parent.setColor(Color.RED);
                    node = node.parent.parent;
                } else {
                    // Case 2: 节点是父节点的右孩子
                    if (node == node.parent.right) {
                        node = node.parent;
                        leftRotate(node);
                    }
                    // Case 3: 节点是父节点的左孩子
                    node.parent.setColor(Color.BLACK);
                    node.parent.parent.setColor(Color.RED);
                    rightRotate(node.parent.parent);
                }
            } else {
                // 对称情况（父节点是右孩子）
                RBNode<T> uncle = node.parent.parent.left;

                if (uncle.getColor() == Color.RED) {
                    node.parent.setColor(Color.BLACK);
                    uncle.setColor(Color.BLACK);
                    node.parent.parent.setColor(Color.RED);
                    node = node.parent.parent;
                } else {
                    if (node == node.parent.left) {
                        node = node.parent;
                        rightRotate(node);
                    }
                    node.parent.setColor(Color.BLACK);
                    node.parent.parent.setColor(Color.RED);
                    leftRotate(node.parent.parent);
                }
            }
        }
        root.setColor(Color.BLACK); // 确保根节点为黑色
    }

    // 左旋操作
    private void leftRotate(RBNode<T> x) {
        RBNode<T> y = x.right;
        x.right = y.left;

        if (y.left != NIL) {
            y.left.parent = x;
        }

        y.parent = x.parent;

        if (x.parent == NIL) {
            root = y;
        } else if (x == x.parent.left) {
            x.parent.left = y;
        } else {
            x.parent.right = y;
        }

        y.left = x;
        x.parent = y;
    }

    // 右旋操作
    private void rightRotate(RBNode<T> y) {
        RBNode<T> x = y.left;
        y.left = x.right;

        if (x.right != NIL) {
            x.right.parent = y;
        }

        x.parent = y.parent;

        if (y.parent == NIL) {
            root = x;
        } else if (y == y.parent.left) {
            y.parent.left = x;
        } else {
            y.parent.right = x;
        }

        x.right = y;
        y.parent = x;
    }

    // 中序遍历（验证排序性质）
    public void inOrder() {
        inOrderTraversal(root);
    }

    private void inOrderTraversal(RBNode<T> node) {
        if (node != NIL) {
            inOrderTraversal(node.left);
            System.out.print(node.getData() + "(" + node.getColor() + ") ");
            inOrderTraversal(node.right);
        }
    }

    // 测试用例
    public static void main(String[] args) {
        RedBlackTree<Integer> tree = new RedBlackTree<>();

        // 插入测试数据
        int[] values = {7, 3, 18, 10, 22, 8, 11, 26};
        for (int val : values) {
            tree.insert(val);
        }

        System.out.println("中序遍历结果:");
        tree.inOrder(); // 输出: 3(BLACK) 7(RED) 8(BLACK) 10(RED) 11(BLACK) 18(BLACK) 22(RED) 26(RED)
    }

}
